"""
MCP SSH服务器实现
提供SSH连接和命令执行功能供LLM调用
"""

import json
import os
import paramiko
from typing import Any, Dict, List, Optional

# 配置文件路径
CONFIG_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'config', 'servers.json')

class SSHServer:
    """SSH服务器管理类"""
    
    def __init__(self):
        self.ensure_config_dir()
    
    def ensure_config_dir(self):
        """确保配置目录存在"""
        config_dir = os.path.dirname(CONFIG_PATH)
        if not os.path.exists(config_dir):
            os.makedirs(config_dir, recursive=True)
        
        # 如果配置文件不存在，创建默认配置
        if not os.path.exists(CONFIG_PATH):
            default_config = {"servers": []}
            with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
                json.dump(default_config, f, ensure_ascii=False, indent=2)
    
    def read_config(self) -> Dict[str, Any]:
        """读取配置文件"""
        try:
            with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
                return json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            return {"servers": []}
    
    def save_config(self, config: Dict[str, Any]) -> bool:
        """保存配置文件"""
        try:
            with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存配置失败: {e}")
            return False
    
    def list_servers(self) -> List[Dict[str, Any]]:
        """列出所有服务器"""
        config = self.read_config()
        return config.get("servers", [])
    
    def add_server(self, **kwargs) -> Dict[str, Any]:
        """添加新服务器"""
        config = self.read_config()
        servers = config.get("servers", [])
        
        # 检查服务器名称是否已存在
        if any(server["name"] == kwargs["name"] for server in servers):
            return {"success": False, "error": "服务器名称已存在"}
        
        # 创建新服务器配置
        new_server = {
            "name": kwargs["name"],
            "ip": kwargs["ip"],
            "port": kwargs.get("port", 22),
            "username": kwargs["username"],
            "auth_type": kwargs["auth_type"],
            "description": kwargs.get("description", "")
        }
        
        # 根据认证方式添加相应字段
        if kwargs["auth_type"] == "password":
            new_server["password"] = kwargs.get("password", "")
        elif kwargs["auth_type"] == "key":
            new_server["key_path"] = kwargs.get("key_path", "")
        
        servers.append(new_server)
        config["servers"] = servers
        
        if self.save_config(config):
            return {"success": True, "message": "服务器添加成功"}
        else:
            return {"success": False, "error": "保存配置失败"}
    
    def delete_server(self, server_name: str) -> Dict[str, Any]:
        """删除服务器"""
        config = self.read_config()
        servers = config.get("servers", [])
        
        # 查找并删除服务器
        original_count = len(servers)
        servers = [server for server in servers if server["name"] != server_name]
        
        if len(servers) == original_count:
            return {"success": False, "error": "服务器不存在"}
        
        config["servers"] = servers
        
        if self.save_config(config):
            return {"success": True, "message": "服务器删除成功"}
        else:
            return {"success": False, "error": "保存配置失败"}
    
    def get_server_info(self, server_name: str) -> Dict[str, Any]:
        """获取服务器信息"""
        config = self.read_config()
        servers = config.get("servers", [])
        
        for server in servers:
            if server["name"] == server_name:
                # 隐藏敏感信息
                safe_server = server.copy()
                if "password" in safe_server:
                    safe_server["password"] = "******"
                return {"success": True, "server": safe_server}
        
        return {"success": False, "error": "服务器不存在"}
    
    def execute_command(self, server_name: str, command: str) -> Dict[str, Any]:
        """在指定服务器上执行SSH命令"""
        try:
            # 获取服务器信息
            server_info = self.get_server_info(server_name)
            if not server_info["success"]:
                return server_info
            
            server = server_info["server"]
            
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接参数
            connect_kwargs = {
                "hostname": server["ip"],
                "port": server["port"],
                "username": server["username"],
                "timeout": 10
            }
            
            # 根据认证方式设置连接参数
            if server["auth_type"] == "password":
                if not server.get("password"):
                    return {"success": False, "error": "密码未配置"}
                connect_kwargs["password"] = server["password"]
            elif server["auth_type"] == "key":
                if not server.get("key_path"):
                    return {"success": False, "error": "密钥路径未配置"}
                connect_kwargs["key_filename"] = server["key_path"]
            
            # 建立连接
            ssh.connect(**connect_kwargs)
            
            # 执行命令
            stdin, stdout, stderr = ssh.exec_command(command)
            
            # 获取输出
            output = stdout.read().decode('utf-8', errors='replace')
            error = stderr.read().decode('utf-8', errors='replace')
            
            # 关闭连接
            ssh.close()
            
            return {
                "success": True,
                "output": output,
                "error": error if error else None
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def test_connection(self, server_name: str) -> Dict[str, Any]:
        """测试与指定服务器的SSH连接"""
        try:
            # 获取服务器信息
            server_info = self.get_server_info(server_name)
            if not server_info["success"]:
                return server_info
            
            server = server_info["server"]
            
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接参数
            connect_kwargs = {
                "hostname": server["ip"],
                "port": server["port"],
                "username": server["username"],
                "timeout": 10
            }
            
            # 根据认证方式设置连接参数
            if server["auth_type"] == "password":
                if not server.get("password"):
                    return {"success": False, "error": "密码未配置"}
                connect_kwargs["password"] = server["password"]
            elif server["auth_type"] == "key":
                if not server.get("key_path"):
                    return {"success": False, "error": "密钥路径未配置"}
                connect_kwargs["key_filename"] = server["key_path"]
            
            # 建立连接
            ssh.connect(**connect_kwargs)
            
            # 关闭连接
            ssh.close()
            
            return {
                "success": True,
                "message": f"与服务器 {server_name} 的连接测试成功"
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)} 

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("用法: python ssh_server.py <command> <server_name> [command_text]")
        print("可用命令:")
        print("  test_connection <server_name>  - 测试服务器连接")
        print("  execute_command <server_name> <command>  - 执行SSH命令")
        sys.exit(1)
    
    command = sys.argv[1]
    server_name = sys.argv[2]
    
    ssh_server = SSHServer()
    
    try:
        if command == "test_connection":
            result = ssh_server.test_connection(server_name)
        elif command == "execute_command":
            if len(sys.argv) < 4:
                print("错误: execute_command 需要指定要执行的命令")
                sys.exit(1)
            command_text = sys.argv[3]
            result = ssh_server.execute_command(server_name, command_text)
        else:
            print(f"错误: 未知命令 '{command}'")
            sys.exit(1)
        
        # 输出JSON结果
        print(json.dumps(result, ensure_ascii=False, indent=2))
        
        # 根据结果设置退出码
        if result.get("success"):
            sys.exit(0)
        else:
            sys.exit(1)
            
    except Exception as e:
        error_result = {"success": False, "error": str(e)}
        print(json.dumps(error_result, ensure_ascii=False, indent=2))
        sys.exit(1) 